import re, sys, struct, os

regex=re.compile('[\deE\.\-\+]+')

def parse_line(line):
    items=regex.findall(line)
    (x,y,z,t,f)=items
    return (int(x),int(y),int(z),int(t),float(f))
    

def convert(input_filename):
    print 'checking lattice size....'
    input=open(input_filename,'r')
    (nx,ny,nz,nt)=(0,0,0,0)
    for line in input:
        (x,y,z,t,f) = parse_line(line)
        nx=max(nx,x)
        ny=max(ny,y)
        nz=max(nz,z)
        nt=max(nt,t)
    (nx,ny,nz,nt)=(nx+1,ny+1,nz+1,nt+1)
    print 'lattice size: %sx%sx%sx%s' % (nx,ny,nz,nt)
    print 'reading data...'
    input=open(input_filename,'r')
    output=open(input_filename+'.tmp','wb')
    for line in input:
        (x,y,z,t,f) = parse_line(line)
        output.seek((z+nz*(y+ny*(x+nx*t)))*4)
        output.write(struct.pack('>f',f))
    print 'creating vtk file %s ...'
    input=open(input_filename+'.tmp','r')
    output=open(input_filename+'.vtk','wb')
    output.write(    
        "# vtk DataFile Version 2.0\n"
        "%s\n"
        "BINARY\n"
        "DATASET STRUCTURED_POINTS\n"
        "DIMENSIONS %i %i %i\n"
        "ORIGIN     0   0   0\n"
        "SPACING    1   1   1\n"
        "POINT_DATA %i" % (input_filename,nx,ny,nz,nx*ny*nz))
    for t in range(nt):
        data=input.read(nx*ny*nz*4)
        output.write("\nSCALARS scalars_t%i float\nLOOKUP_TABLE default\n" % (t))
        output.write(data)
    os.unlink(input_filename+'.tmp')
    print 'done!'

if __name__=='__main__':
    try:
        convert(sys.argv[1])
    except Exception,e:
        print """
failure to convert (%s)
Usage: scalar2vtk.py filename.dat
where filename.dat is something like
0 0 0 0 1.420671e-01
1 0 0 0 -1.229045e+00
2 0 0 0 -4.986515e-01
...
15 15 15 7 -2.926175e-01
"""  % e
